#include <functional>
#include <tuple>
#include <type_traits>

//xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx 基础设施
// ---------- 轻量级 Monad 概念 ----------
template <template <class...> class M>
concept MonadConcept = requires {
    typename M<int>;  // 至少能实例化
};

// 简单的 Identity Monad，用作“内层”默认 Monad
template <class A>
struct Identity
{
    A value;
    auto runIdentity() const { return value; }
};

//xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
template <class W, class A, template <class...> class Inner = Identity>
struct WriterT;

template <class R, class A, template <class...> class Inner = Identity>
struct ReaderT;

template <class S, class A, template <class...> class Inner = Identity>
struct StateT;

//xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
template <class A, class F>
auto operator>>=(Identity<A> const& m, F&& f)
{
    return std::invoke(std::forward<F>(f), m.value);
}
//
auto mreturn = [](auto x) { return Identity{std::move(x)}; };

//xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx Writer Monad / WriterT
template <class W, class A, template <class...> class Inner = Identity>
struct WriterT
{
    using InnerType = Inner<std::tuple<A, W>>;
    InnerType runWriterT;

    template <class B>
    friend WriterT<W, B, Inner> operator>>(WriterT<W, A, Inner> const& lhs, WriterT<W, B, Inner> const& rhs)
    {
        return lhs >>= [rhs](auto&&) { return rhs; };
    }

    static constexpr auto tell(W w) { return WriterT<W, std::monostate, Inner>{mreturn(std::make_tuple(std::monostate{}, w))}; }
};

// 纯 Writer 是 WriterT 特例
template <class W, class A>
using Writer = WriterT<W, A, Identity>;

// 运行函数
template <class W, class A, template <class...> class I>
auto runWriter(WriterT<W, A, I> const& w)
{
    return w.runWriterT;
}

// >>= 实现
template <class W, class A, class F, template <class...> class I>
auto operator>>=(WriterT<W, A, I> const& m, F&& f)
{
    return WriterT<W, std::invoke_result_t<F, A>, I>{m.runWriterT >>= [&](auto const& pair)
                                                     {
                                                         auto [a, w1] = pair;
                                                         auto w2 = std::invoke(std::forward<F>(f), a);
                                                         return w2.runWriterT >>= [=](auto const& p2)
                                                         {
                                                             auto [b, w2v] = p2;
                                                             return mreturn(std::make_tuple(b, w1 + w2v));  // 要求 W 支持 +
                                                         };
                                                     }};
}

//xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx Reader Monad / ReaderT
// ---------- Reader ----------
template <class R, class A, template <class...> class Inner = Identity>
struct ReaderT
{
    using InnerType = Inner<A(R)>;
    std::function<Inner<A>(R)> runReaderT;

    template <class B>
    friend ReaderT<R, B, Inner> operator>>(ReaderT<R, A, Inner> const& lhs, ReaderT<R, B, Inner> const& rhs)
    {
        return lhs >>= [rhs](auto&&) { return rhs; };
    }

    static constexpr auto ask()
    {
        return ReaderT<R, R, Inner>{[](R r) { return mreturn(r); }};
    }
    template <class F>
    static constexpr auto local(F f, ReaderT<R, A, Inner> m)
    {
        return ReaderT<R, A, Inner>{[=](R r) { return m.runReaderT(f(r)); }};
    }
};

// 纯 Reader
template <class R, class A>
using Reader = ReaderT<R, A, Identity>;

template <class R, class A, template <class...> class I>
auto runReader(ReaderT<R, A, I> const& r, R env)
{
    return r.runReaderT(env);
}

// >>=
template <class R, class A, class F, template <class...> class I>
auto operator>>=(ReaderT<R, A, I> const& m, F&& f)
{
    return ReaderT<R, std::invoke_result_t<F, A>, I>{
        [=](R r) { return m.runReaderT(r) >>= [&](A a) { return std::invoke(std::forward<F>(f), a).runReaderT(r); }; }};
}

//xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx State Monad / StateT
// S: state; A: obj
template <class S, class A, template <class...> class Inner = Identity>
struct StateT
{
    using InnerType = Inner<std::tuple<A, S>(S)>;
    std::function<Inner<std::tuple<A, S>>(S)> runStateT;

    // operator >>
    template <class B>
    friend StateT<S, B, Inner> operator>>(StateT<S, A, Inner> const& lhs, StateT<S, B, Inner> const& rhs)
    {
        return lhs >>= [rhs](auto&&) { return rhs; };
    }

    static constexpr auto get()
    {
        return StateT<S, S, Inner>{[](S s) { return mreturn(std::make_tuple(s, s)); }};
    }
    static constexpr auto put(S s)
    {
        return StateT<S, std::monostate, Inner>{[=](S) { return mreturn(std::make_tuple(std::monostate{}, s)); }};
    }
};

// 纯 State
template <class S, class A>
using State = StateT<S, A, Identity>;

template <class S, class A, template <class...> class I>
auto runState(StateT<S, A, I> const& s, S st)
{
    return s.runStateT(st);
}

// >>=
template <class S, class A, class F, template <class...> class I>
auto operator>>=(StateT<S, A, I> const& m, F&& f)
{
    return StateT<S, std::invoke_result_t<F, A>, I>{[=](S s)
                                                    {
                                                        return m.runStateT(s) >>= [&](auto const& pair)
                                                        {
                                                            auto [a, s1] = pair;
                                                            return std::invoke(std::forward<F>(f), a).runStateT(s1);
                                                        };
                                                    }};
}
